#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
#include "ffmpeg_log.h"

int decode_video(const char *videoFilePath);

int encode_picture(AVCodecContext *codecContext, const char *picturePath, AVFrame *frameYUV);

int main() {
    const char *videoFilePath = "D:\\ffmpeg-demo\\可爱宝宝.mp4";
    decode_video(videoFilePath);
    return 0;
}

int decode_video(const char *videoFilePath) {
    const char *pictureDirPath = "D:\\ffmpeg-demo\\video_picture\\可爱宝宝";
    AVFormatContext *formatContext = avformat_alloc_context();
    int ret;
    if ((ret = avformat_open_input(&formatContext, videoFilePath, NULL, NULL)) < 0) {
        printfError(ret, "open_input");
        return ret;
    }
    if ((ret = avformat_find_stream_info(formatContext, NULL)) < 0) {
        printfError(ret, "find_stream");
        return ret;
    }
    int videoIndex = -1;
    for (int i = 0; i < formatContext->nb_streams; ++i) {
        if (formatContext->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            videoIndex = i;
            break;
        }
    }
    if (videoIndex < 0) {
        printf("file in not a video file.\n");
        return -1;
    }
    AVStream *stream = formatContext->streams[videoIndex];
    AVCodec *codec = avcodec_find_decoder(stream->codecpar->codec_id);
    AVCodecContext *codecContext = avcodec_alloc_context3(codec);
    avcodec_parameters_to_context(codecContext, stream->codecpar);
    if ((ret = avcodec_open2(codecContext, codec, NULL)) < 0) {
        printfError(ret, "open codec");
        return ret;
    }

    AVFrame *frame = av_frame_alloc();
    AVFrame *frameYuv = av_frame_alloc();
    AVPacket *packet = av_packet_alloc();
    struct SwsContext *swsContext = sws_getContext(codecContext->width, codecContext->height, codecContext->pix_fmt,
                                                   codecContext->width, codecContext->height, AV_PIX_FMT_YUV420P,
                                                   SWS_BICUBIC, NULL, NULL, NULL);
    unsigned char *pictureBuff = av_malloc(
            av_image_get_buffer_size(AV_PIX_FMT_YUV420P, codecContext->width, codecContext->height, 1));
    av_image_fill_arrays(frameYuv->data, frameYuv->linesize, pictureBuff, AV_PIX_FMT_YUV420P, codecContext->width,
                         codecContext->height, 1);
    while (1) {
        ret = av_read_frame(formatContext, packet);
        if (ret < 0) {
            printfError(ret, "read_frame");
            return ret;
        }
        if (packet->stream_index != videoIndex) {
            continue;
        }
        ret = avcodec_send_packet(codecContext, packet);
        if (ret < 0) {
            printfError(ret, "send packet");
            continue;
        }
        ret = avcodec_receive_frame(codecContext, frame);
        if (ret < 0) {
            printfError(ret, "receive_frame");
            continue;
        }
        if (frame->pict_type != AV_PICTURE_TYPE_I) {
            continue;
        }
        log_packet(formatContext, packet, "encode");
        ret = sws_scale(swsContext, (const uint8_t *const *) frame->data, frame->linesize, 0, codecContext->height,
                        frameYuv->data, frameYuv->linesize);
        if (ret < 0) {
            printfError(ret, "sws scale");
            continue;
        }
        const char *ptsStr = av_ts2str(packet->pts);
        int strLength = strlen(pictureDirPath) + strlen(ptsStr) + 5;
        char *picturePath = malloc(strLength * sizeof(char));
        sprintf(picturePath, "%s_%s.jpg", pictureDirPath, ptsStr);
        encode_picture(codecContext, picturePath, frameYuv);
        av_packet_unref(packet);
    }
}

int encode_picture(AVCodecContext *codecContext, const char *picturePath, AVFrame *frameYUV) {
    AVFormatContext *picFormatContext = avformat_alloc_context();
    AVOutputFormat *outputFormat = av_guess_format("mjpeg", NULL, NULL);
    picFormatContext->oformat = outputFormat;
    if (avio_open(&picFormatContext->pb, picturePath, AVIO_FLAG_READ_WRITE) < 0) {
        printf("open output file error.\n");
        return -1;
    }
    AVStream *pictureStream = avformat_new_stream(picFormatContext, 0);
    if (pictureStream == NULL) {
        return -1;
    }
    AVCodecContext *picCodecContext = pictureStream->codec;
    picCodecContext->codec_id = outputFormat->video_codec;
    picCodecContext->codec_type = AVMEDIA_TYPE_VIDEO;
    picCodecContext->pix_fmt = AV_PIX_FMT_YUVJ420P;
    picCodecContext->width = codecContext->width;
    picCodecContext->height = codecContext->height;
    picCodecContext->time_base.num = 1;
    picCodecContext->time_base.den = 25;
    av_dump_format(picFormatContext, 0, picturePath, 1);
    AVCodec *picCodec = avcodec_find_encoder(picCodecContext->codec_id);
    if (!picCodec) {
        printf("codec not found.\n");
        return -1;
    }
    if (avcodec_open2(picCodecContext, picCodec, NULL) < 0) {
        printf("open codec error.\n");
        return -1;
    }
    AVFrame *picture = av_frame_alloc();
    int size = av_image_get_buffer_size(picCodecContext->pix_fmt, picCodecContext->width,
                                        picCodecContext->height, 1);
    uint8_t *pictureBuff = av_malloc(size);
    if (!pictureBuff) {
        printf("create picture buff error.\n");
        return -1;
    }
    av_image_fill_arrays(picture->data, picture->linesize, pictureBuff, picCodecContext->pix_fmt,
                         picCodecContext->width,
                         picCodecContext->height, 1);
    avformat_write_header(picFormatContext, NULL);
    picture->data[0] = frameYUV->data[0];
    picture->data[1] = frameYUV->data[1];
    picture->data[2] = frameYUV->data[2];
    int y_size = picCodecContext->width * picCodecContext->height;
    AVPacket *packet = av_malloc(sizeof(AVPacket));;
    av_new_packet(packet, y_size * 3);
    int got_picture;
    int ret = avcodec_encode_video2(picCodecContext, packet, picture, &got_picture);
    if (ret < 0) {
        printf("encode error.\n");
        return -1;
    }
    if (got_picture == 1) {
        packet->stream_index = pictureStream->index;
        av_write_frame(picFormatContext, packet);
        av_write_trailer(picFormatContext);
        printf("encode successful.\n");
    } else {
        return -1;
    }
    avcodec_close(pictureStream->codec);
    av_free(pictureBuff);
    av_frame_free(&picture);
    avcodec_close(picCodecContext);
    avformat_close_input(&picFormatContext);
    return 0;
}
